from typing import Callable

from zkl_aiutils_training import SteppingTaskPlugin

from .metrics import MetricsCollector, get_metric_scalar_value


class EarlyStopByMetric(SteppingTaskPlugin, MetricsCollector):
    def __init__(self, metric_name: str, criteria: Callable[[int | float], bool] | None):
        self.metric_name = metric_name
        self.criteria = criteria

    def on_after_step(self):
        if self._check_should_stop():
            raise StopIteration

    def _check_should_stop(self):
        if self.criteria is None:
            return False
        value = self._get_metric_value(self.metric_name)
        if value is None:
            return False
        value = get_metric_scalar_value(value)
        if value is None:
            raise ValueError(f"Metric {self.metric_name} is not a scalar!")
        return self.criteria(value)


class EarlyStopByMaxThreshold(EarlyStopByMetric):
    def __init__(self, metric_name: str, threshold: int | float | None):
        criteria = (lambda value: value >= threshold) if threshold is not None else None
        super().__init__(metric_name, criteria)


class EarlyStopByMinThreshold(EarlyStopByMetric):
    def __init__(self, metric_name: str, threshold: int | float | None):
        criteria = (lambda value: value <= threshold) if threshold is not None else None
        super().__init__(metric_name, criteria)
